#include <ros/ros.h>
#include <mavros_msgs/ActuatorControl.h>
#include <mavros_msgs/CommandBool.h>
#include <mavros_msgs/SetMode.h>
#include <mavros_msgs/State.h>
#include <mavros_msgs/Altitude.h>
#include <sensor_msgs/BatteryState.h>
#include "robotiq_ft_sensor/ft_sensor.h"
#include <fstream>
#include <string>

mavros_msgs::ActuatorControl motor_control_msg;
mavros_msgs::State current_state;
int current_pwm = 0;
double force_z = 0.0;
double moment_z = 0.0;
double battery_voltage = 0.0;
double battery_current = 0.0;
bool torque_sensor_avaliable = false;

void stateCallback(const mavros_msgs::State::ConstPtr& msg){
    current_state = *msg;
}

void pwmCallback(const mavros_msgs::Altitude::ConstPtr& msg){
    current_pwm = (int)msg->monotonic;   /// Motor test use this msg to transmit pwm
}

void forceTorqueCallback(const robotiq_ft_sensor::ft_sensor &msg){
    if(msg.Fz == msg.Fz && fabs(msg.Fz < 10000.f)){
        force_z = msg.Fz;
        torque_sensor_avaliable = true;
    }else{
        ROS_WARN("Received an invalid force value.");
    }

    if(msg.Mz == msg.Mz && fabs(msg.Mz < 10000.f)){
        moment_z = msg.Mz;
    } else{
        ROS_WARN("Received an invalid moment value.");
    }
}

void batteryCallback(const sensor_msgs::BatteryState::ConstPtr& msg){
    battery_voltage = msg->voltage;
    battery_current = msg->current;
}

int main(int argc, char** argv){
    ros::init(argc, argv, "motor_control_eg");
    ros::NodeHandle nh;

    std::string save_path;
    ros::param::get("/motor_mavlink_control/save_path", save_path);
    std::cout << "Test result will be saved in: " << save_path <<std::endl;

    ros::Subscriber state_sub = nh.subscribe("/mavros/state", 1, stateCallback);
    ros::Subscriber pwm_sub = nh.subscribe("/mavros/altitude", 1, pwmCallback);   /// Motor test use this msg to transmit pwm
    ros::Subscriber force_sub = nh.subscribe("/robotiq_ft_sensor", 1, forceTorqueCallback);
    ros::Subscriber battery_sub = nh.subscribe("/mavros/battery", 1, batteryCallback);

    ros::Publisher motor_control_pub = nh.advertise<mavros_msgs::ActuatorControl>("/mavros/actuator_control", 1, true);
    ros::ServiceClient arming_client = nh.serviceClient<mavros_msgs::CommandBool>
            ("mavros/cmd/arming");
    ros::ServiceClient set_mode_client = nh.serviceClient<mavros_msgs::SetMode>
            ("mavros/set_mode");

    ROS_WARN("Motor Test started!");
    // wait for FCU connection
    ros::Rate wait_rate(20.0);
    while(current_state.connected){
        ros::spinOnce();
        wait_rate.sleep();
    }
    ROS_WARN("FCU connected");

    // wait for torque sensor connection
    while(!torque_sensor_avaliable){
        ros::spinOnce();
        wait_rate.sleep();
    }
    ROS_WARN("Torque sensor connected");

    // Collect zero point for torque sensor
    double force_z_zero = 0.0;
    double moment_z_zero = 0.0;
    int fm_collection_times = 20;
    for(int i=0; i<fm_collection_times; i++){
        force_z_zero += force_z;
        moment_z_zero += moment_z;

        ros::spinOnce();
        wait_rate.sleep();
    }
    force_z_zero = force_z_zero / fm_collection_times;
    moment_z_zero = moment_z_zero / fm_collection_times;

    // Set init control value
    const int control_vector_size = 8;
    float controls[control_vector_size] = {-1.f, -1.f, -1.f, -1.f, 0.f, 0.f, 0.f, 0.f};

    ROS_WARN("Starting offboard...");

    // Change mode
    mavros_msgs::SetMode offb_set_mode;
    offb_set_mode.request.custom_mode = "OFFBOARD";
    mavros_msgs::CommandBool arm_cmd;
    arm_cmd.request.value = true;

    // Precheck
    int status_counter = 0;
    while(ros::ok())
    {
        //send a few setpoints every time
        motor_control_msg.header.stamp = ros::Time::now();
        motor_control_msg.group_mix = 0;
        for(int i=0; i<control_vector_size; i++){
            motor_control_msg.controls[i] = controls[i];
        }
        motor_control_pub.publish(motor_control_msg);

        //Set mode by service
        if(status_counter == 100){
            if( set_mode_client.call(offb_set_mode) &&
                offb_set_mode.response.mode_sent){
                ROS_WARN("Offboard enabled");
            }
        }
        else if(status_counter == 200){
            if( arming_client.call(arm_cmd) &&
                arm_cmd.response.success){
                ROS_WARN("Vehicle armed");
            }
            break;
        }
        status_counter ++;
        ros::spinOnce();
        wait_rate.sleep();
    }

    ROS_WARN("Caution: Vehicle is going to fly! :) ");

    /********* Test Moment and Force with different PWM *********/
    const int loop_frequency = 100;
    ros::Rate loop_rate(loop_frequency);

    std::ofstream outFile1;
    std::string file1 = save_path + "force_moment_pwm_test.csv";
    outFile1.open(file1, std::ios::out);
    outFile1 << "Input" <<',' << "PWM" <<',' << "Force" << ',' << "Moment" <<',' << "Voltage" << ',' << "Current" << '\n';

    for(int i=0; i<control_vector_size; i++){controls[i] = -1.f;} //initialize
    int counter_times_each = 200;
    int counter_times_stable = 100;
    int counter_force_test = 0;
    int counter_force_test_each = 0;
    double avg_force = 0.0;
    double avg_moment = 0.0;

    while (ros::ok()){
        if(counter_force_test == counter_times_stable) {
            counter_force_test_each = 0;
            avg_force = 0.0;
            avg_moment = 0.0;
        }else if(counter_force_test == counter_times_each){
            //One record finished
            avg_force = avg_force / counter_force_test_each;
            avg_moment = avg_moment / counter_force_test_each;
            outFile1 << controls[0] <<',' << current_pwm <<','<< avg_force <<','<< avg_moment <<','<< battery_voltage <<','<< battery_current <<'\n';
            for(int i=0; i<control_vector_size; i++){controls[i] += 0.1f;}

            if(controls[0] > 1.01f){
                ROS_WARN("Force-Moment-PWM Test finished");
                break;
            }
            counter_force_test = 0;

        }else if(counter_force_test > counter_times_stable){
            counter_force_test_each ++;
            avg_force += force_z - force_z_zero;
            avg_moment += moment_z - moment_z_zero;
        }

        motor_control_msg.header.stamp = ros::Time::now();
        motor_control_msg.group_mix = 0;
        for(int i=0; i<control_vector_size; i++){ motor_control_msg.controls[i] = controls[i];}
        motor_control_pub.publish(motor_control_msg);

        counter_force_test ++;

        ros::spinOnce();
        loop_rate.sleep();
    }

    outFile1.close();

    /**** Pause a while ****/
    int pose_time_counter = 0;
    const int pause_times = 1000;
    while(ros::ok()){
        motor_control_msg.header.stamp = ros::Time::now();
        motor_control_msg.group_mix = 0;
        for(int i=0; i<control_vector_size; i++){ motor_control_msg.controls[i] = -1.f;}
        motor_control_pub.publish(motor_control_msg);

        pose_time_counter ++;
        if(pose_time_counter > pause_times){
            break;
        }

        ros::spinOnce();
        loop_rate.sleep();
    }

    /******* Test Ramp signal *******/
    std::ofstream outFile2;
    std::string file2 = save_path + "force_moment_ramp_test.csv";
    outFile2.open(file2, std::ios::out);
    outFile2 << "Input" <<',' << "PWM" <<',' << "Force" << ',' << "Moment" <<',' << "Voltage" << ',' << "Current" << '\n';

    // Set testing PWM
    const float start_value = -1.f;
    const float end_vlaue = 1.f;
    float period = 8.f;  // seconds

    float step = (end_vlaue - start_value) / (period / 2) / (float)loop_frequency;
    int timer_counter = 0;
    for(int i=0; i<control_vector_size; i++){controls[i] = -1.f;} //initialize

    while (ros::ok()){
        // The signal climbs up and then get down and repeat.
        if(timer_counter < loop_frequency*period/2){
            for(int i=0; i<4; i++){controls[i] += step;}  //Only change channel 1 - 4
        }else if(timer_counter < loop_frequency*period){
            for(int i=0; i<4; i++){controls[i] -= step;}
        }else{
            timer_counter = 0;
            ROS_WARN("One loop finished!");
            period -= 2.f;
            for(int i=0; i<control_vector_size; i++){controls[i] = -1.f;} //initialize
            if(period < 1.f){
                break;
            }
        }
        timer_counter ++;
//        for(int i=0; i<control_vector_size; i++){controls[i] += 0.1f;}

        // Set values
        motor_control_msg.header.stamp = ros::Time::now();
        motor_control_msg.group_mix = 0;
        for(int i=0; i<control_vector_size; i++){ motor_control_msg.controls[i] = controls[i];}
        motor_control_pub.publish(motor_control_msg);

        outFile2 << controls[0] <<',' << current_pwm <<','<< force_z-force_z_zero <<','<< moment_z-moment_z_zero <<','<< battery_voltage <<','<< battery_current <<'\n';

        ros::spinOnce();
        loop_rate.sleep();
    }
    outFile2.close();

    /*** Test end. Shutdown motor ****/
    for(int i=0; i<control_vector_size; i++){controls[i] = -1.f;} //initialize
    while(ros::ok()){
        motor_control_msg.header.stamp = ros::Time::now();
        motor_control_msg.group_mix = 0;
        for(int i=0; i<control_vector_size; i++){ motor_control_msg.controls[i] = controls[i];}
        motor_control_pub.publish(motor_control_msg);

        ROS_WARN_THROTTLE(1.0, " Test finished. You may disconnect the battery now.");
        if(!current_state.connected){
            break;
        }

        ros::spinOnce();
        loop_rate.sleep();
    }

    return 0;
}

